Gemma/[Gemma_2]Keras_Quickstart.ipynb

{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "Tce3stUlHN0L" }, "source": [ "##### Copyright 2025 Google LLC." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "tuOe1ymfHZPu" }, "outputs": [], "source": [ "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "PXNm5_p_oxMF" }, "source": [ "This is a quick demo of Gemma running on KerasNLP.\n", "\n", "This notebook will run on an NVIDIA T4 GPU.\n", "\n", "General Keras reading:\n", "- [Getting started with Keras](https://keras.io/getting_started/)\n", "- [Getting started with KerasNLP](https://keras.io/guides/keras_nlp/getting_started/)\n", "\n", "<table align=\"left\">\n", " <td>\n", " <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Gemma/[Gemma_2]Keras_Quickstart.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n", " </td>\n", "</table>" ] }, { "cell_type": "markdown", "metadata": { "id": "VzFhJDayXroT" }, "source": [ "## Access" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "DrBoa_Urw9Vx" }, "outputs": [], "source": [ "import os\n", "from google.colab import userdata\n", "\n", "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n", "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')\n", "\n", "os.environ[\"KERAS_BACKEND\"] = \"jax\" # Or \"tensorflow\" or \"torch\"." ] }, { "cell_type": "markdown", "metadata": { "id": "z9oy3QUmXtSd" }, "source": [ "## Installation" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "id": "bMboT70Xop8G" }, "outputs": [], "source": [ "%%capture\n", "\n", "# install all dependencies\n", "!pip install -U keras-nlp\n", "!pip install -U keras>=3.3.3" ] }, { "cell_type": "markdown", "metadata": { "id": "FX47AUYrXwLK" }, "source": [ "## Quickstart" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "id": "ww83zI9ToPso" }, "outputs": [ { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Preprocessor: \"gemma_causal_lm_preprocessor\"\n", "</pre>\n" ], "text/plain": [ "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", "┃ Layer (type) ┃ Config ┃\n", "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", "│ gemma_tokenizer (GemmaTokenizer) │ Vocab size: 256,000 │\n", "└───────────────────────────────────────────────────────────────┴──────────────────────────────────────────┘\n", "</pre>\n" ], "text/plain": [ "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", "┃\u001b[1m \u001b[0m\u001b[1mLayer (type) \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m Config\u001b[0m\u001b[1m \u001b[0m┃\n", "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m) │ Vocab size: \u001b[38;5;34m256,000\u001b[0m │\n", "└───────────────────────────────────────────────────────────────┴──────────────────────────────────────────┘\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Model: \"gemma_causal_lm\"\n", "</pre>\n" ], "text/plain": [ "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", "┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃\n", "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", "│ padding_mask (InputLayer) │ (None, None) │ 0 │ - │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ token_ids (InputLayer) │ (None, None) │ 0 │ - │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ gemma_backbone │ (None, None, 2304) │ 2,614,341,888 │ padding_mask[0][0], │\n", "│ (GemmaBackbone) │ │ │ token_ids[0][0] │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ token_embedding │ (None, None, 256000) │ 589,824,000 │ gemma_backbone[0][0] │\n", "│ (ReversibleEmbedding) │ │ │ │\n", "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n", "</pre>\n" ], "text/plain": [ "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", "┃\u001b[1m \u001b[0m\u001b[1mLayer (type) \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to \u001b[0m\u001b[1m \u001b[0m┃\n", "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m) │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m) │ \u001b[38;5;34m0\u001b[0m │ - │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m) │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m) │ \u001b[38;5;34m0\u001b[0m │ - │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ gemma_backbone │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2304\u001b[0m) │ \u001b[38;5;34m2,614,341,888\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m], │\n", "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m) │ │ │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m] │\n", "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n", "│ token_embedding │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m) │ \u001b[38;5;34m589,824,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m] │\n", "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m) │ │ │ │\n", "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"> Total params: 2,614,341,888 (4.87 GB)\n", "</pre>\n" ], "text/plain": [ "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,614,341,888\u001b[0m (4.87 GB)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"> Trainable params: 2,614,341,888 (4.87 GB)\n", "</pre>\n" ], "text/plain": [ "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m2,614,341,888\u001b[0m (4.87 GB)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"> Non-trainable params: 0 (0.00 B)\n", "</pre>\n" ], "text/plain": [ "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "import keras_nlp\n", "import keras\n", "\n", "# run at half precision\n", "keras.config.set_floatx(\"bfloat16\")\n", "\n", "# use 2B pre-trained checkpoint\n", "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma2_2b_en\")\n", "gemma_lm.summary()" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "id": "aae5GHrdpj2_" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "It was a dark and stormy night.\n", "\n", "The wind howled and the rain pelted the windows of the house.\n", "\n", "The storm was so bad that the power went out.\n", "\n", "The family was huddled around the fireplace, trying to stay warm.\n", "\n", "Suddenly, there was a knock at the door.\n", "\n", "The family was surprised to see a stranger standing there.\n", "\n", "The stranger said, \"I'm sorry to bother you, but I'm looking for a place to stay.\n", "\n", "The family was hesitant to let the stranger in, but they finally agreed.\n", "\n", "The stranger was a young man named Jack.\n", "\n", "He had been traveling for months, and he was tired and hungry.\n", "\n", "The family offered him some food and a place to sleep.\n", "\n", "Jack was grateful for their kindness, and he stayed with them for a few days.\n", "\n", "One night, Jack was awakened by a loud noise.\n", "\n", "He looked out the window and saw a strange light in the sky.\n", "\n", "He knew that it was a sign of trouble, and he was right.\n", "\n", "The next morning, the family was awakened by the sound of thunder.\n", "\n", "They looked out the window and saw a huge storm cloud approaching.\n", "\n", "The family was scared, but they knew that they had to stay and face the storm.\n", "\n", "The storm was so bad that the family was trapped in the house.\n", "\n", "They were afraid that the storm would destroy their home.\n", "\n", "But then, something amazing happened.\n", "\n", "The storm cloud began to change.\n", "\n", "It became a beautiful rainbow.\n", "\n", "The family was amazed, and they knew that the storm was over.\n", "\n", "The family was safe, and they were grateful for the rainbow.\n", "\n", "They knew that it was a sign of hope and a reminder that even in the darkest of times, there is always a light at the end of the tunnel.\n" ] } ], "source": [ "result = gemma_lm.generate(\"It was a dark and stormy night.\", max_length=512)\n", "print(result)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "name": "[Gemma_2]Keras_Quickstart.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }

Gemma/[Gemma_2]Keras_Quickstart.ipynb (341 lines of code) (raw):